# -*- coding: utf-8 -*-  
'''
nlu相关loss

Created on 2021年10月2日
@author: luoyi
'''
import tensorflow as tf

import utils.conf as conf
from models.crf.process import score, score_full_path


#    nlu相关loss
class NLULoss(tf.keras.losses.Loss):
    '''
        loss = loss_inform + loss_crf
        loss_crf = log(Z(i)) - score(i)        （详见CRF文档）
        loss_inform = 交叉熵
    '''
    def __init__(self, 
                 name='nlu_loss',
                 loss_lamda_crf=conf.NLU.get_loss_lamda_crf(),
                 loss_lamda_inform=conf.NLU.get_loss_lamda_inform(),
                 inform_layer=None,
                 crf_layer=None,
                 **kwargs):
        super(NLULoss, self).__init__(name=name, **kwargs)
        
        self._loss_lamda_crf = loss_lamda_crf
        self._loss_lamda_inform = loss_lamda_inform
        
        self._inform_layer = inform_layer
        self._crf_layer = crf_layer
        pass
    
    def call(self, y_true, y_pred):
        '''
            @param y_true: Tensor(batch_size, 2, max_sen_len)
                            0：意图id。只有第1位有值，其他都是[PAD]-1
                            1：序列标注id。[PAD]=-1
        '''
        #    解析crf预测和意图预测
        pred_inform = self._inform_layer.get_out()
        pred_crf = self._crf_layer.get_out()
        #    解析真实序列和真实意图
        tf.split(y_true, num_or_size_splits=[1,1], axis=1)
        y_inform = y_true[:, 0, :]
        y_crf = y_true[:, 1, :]
        
        #    意图loss
        loss_inform = self.loss_inform(pred_inform, y_inform)
        #    标注loss
        loss_crf = self.loss_crf(pred_crf, y_crf, self._crf_layer.get_transfer())
        
        loss = self._loss_lamda_crf * loss_crf + self._loss_lamda_inform * loss_inform
        
        return loss
    
    #    计算inform损失
    def loss_inform(self, pred_inform, y_inform):
        '''inform损失，交叉熵
            @param pred_inform: 意图分类预测    Tensor(batch_size, inform_size)
            @param y_inform: 真实意图id    Tensor(batch_size, max_sen_len)
        '''
        #    取真实意图id    Tensor(batch_size, )
        y_inform = y_inform[:, 0]
        idx_B = tf.where(y_inform >= 0)[:, 0]
        idx_y_true = tf.stack([idx_B, y_inform], axis=-1)
        
        #    取每个真实意图的预测概率
        pred_inform = tf.gather_nd(params=pred_inform, indices=idx_y_true)
        
        #    计算交叉熵
        loss = -tf.math.log(pred_inform)
        return loss
    
    #    计算crf损失
    def loss_crf(self, pred_crf, y_crf, transfer):
        '''
            @param pred_crf: crf_layer输出    Tensor(batch_size, max_sen_len, pos_size)
            @param y_crf: 真实序列    Tensor(batch_size, max_sen_len)    包括[CLS]和[EOS]的标注
            @param transfer: crf层的状态转移矩阵    Tensor(pos_size, pos_size)
        '''
        score_i = score(pred_crf, y_crf, transfer)
        Z_i = score_full_path(pred_crf, y_crf, transfer)
        loss = Z_i - score_i
        
        return loss
    pass

